from typing import List
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class GraphEntity(BaseModel):
    names: List[str] = Field(description = "文本中的实体")

class GraphRAG:

    def __init__(self, llm_model, graph, prefix_label):
        self.llm_model = llm_model
        self.graph = graph
        self.prefix_label = prefix_label

    def extract_entity(self, query):
        logger.info("GraphRAG extract_entity start")
        template = """
            你是一个专业的信息抽取系统。请从以下文本中精确识别并提取所有有意义的命名实体。

            要求：
            1. 只提取明确提及的实体，不要推测或生成不存在的内容。
            2. 实体类型包括但不限于：人物（Person）、组织（Organization）、地点（Location）、日期/时间（Date）、技术术语（Technology）、产品（Product）、事件（Event）。
            3. 对每个实体，提供实体在文本中的原始形式（保持大小写和标点）
            4. 不要包含任何解释、前缀、后缀或 Markdown。
            
            文本： {context}
        """
        prompt = ChatPromptTemplate.from_template(template)
        extract_chain = prompt | self.llm_model.with_structured_output(GraphEntity)
        extract_result = extract_chain.invoke({ "context": query })
        logger.info(f"GraphRAG extract_entity extract_result: {extract_result}")
        return extract_result.names

    def select_graph(self, entities):
        logger.info("GraphRAG select_graph start")
        logger.info(f"GraphRAG select_graph prefix_label: {self.prefix_label}")
        graph_data = ""
        for entity in entities:
            response = self.graph.query(
                f"""
                    MATCH (n)
                    WHERE any(lbl IN labels(n) WHERE lbl STARTS WITH "{self.prefix_label}")
                    OPTIONAL MATCH (n:`{entity}`)-[r]->(m)
                    RETURN n.id AS source_id, type(r) AS relationship, m.id AS target_id
                    LIMIT 50
                """
            )
            graph_data += "\n".join([f"{ele['source_id']} - {ele['relationship']} - {ele['target_id']}" for ele in response])
        logger.info(f"GraphRAG select_graph graph_data len: {len(graph_data)}")
        return graph_data

    def invoke(self, query):
        logger.info(f"GraphRAG invoke query: {query}")
        entities = self.extract_entity(query)
        graph_data = self.select_graph(entities)

        template = """
            你是一位问题解答专家，请从提供的上下文回答用户问题。
            
            上下文：{context}
            用户问题：{question}
            
            必须用中文详尽回答问题。
        """
        prompt = ChatPromptTemplate.from_template(template)
        generate_chain = prompt | self.llm_model | StrOutputParser()
        response = generate_chain.invoke({ "context": graph_data, "question": query })
        logger.info(f"GraphRAG invoke response len: {len(response)}")
        return { "retrieve_docs": graph_data, "chain_result": response }
